#!/usr/bin/env python3
"""
机器学习模型训练脚本

本脚本提供了完整的机器学习模型训练和测试功能，专门用于智能测试用例生成系统。
支持多种机器学习算法的训练、性能比较、模型保存和测试验证。

主要功能：
- 多模型训练：自动训练朴素贝叶斯、逻辑回归、随机森林三种模型
- 性能评估：提供详细的性能比较和交叉验证
- 最佳模型选择：根据测试准确率自动选择最优模型
- 模型保存：将训练好的模型保存到指定目录
- 模型测试：使用测试用例验证模型性能
- 命令行接口：提供灵活的命令行参数控制

技术特点：
- 自动化训练流程：一键训练所有支持的模型类型
- 性能监控：实时显示训练进度和性能指标
- 数据统计：提供详细的数据集统计信息
- 模型验证：使用多种测试用例验证模型效果
- 灵活配置：支持自定义模型目录和测试集比例

支持的模型类型：
- 朴素贝叶斯（MultinomialNB）：适合文本分类，计算效率高
- 逻辑回归（LogisticRegression）：线性分类器，可解释性强
- 随机森林（RandomForestClassifier）：集成学习，性能稳定


命令行用法：
    # 训练模型
    python train_ml_models.py --train
    
    # 测试模型
    python train_ml_models.py --test models/best_model_random_forest.joblib
    
    # 自定义参数
    python train_ml_models.py --train --model-dir my_models --test-size 0.2

依赖关系：
- ml_trainer：机器学习训练器模块
- ml_dataset：训练数据集模块
- argparse：命令行参数解析
- os, sys：系统操作
"""

# ==================== 导入和配置 ====================

import os
import sys
import argparse
from typing import Optional

# 添加src目录到Python路径
# 确保可以导入src目录下的机器学习模块
sys.path.append('src')

# 导入核心模块
from ml_trainer import MLModelTrainer  # 机器学习训练器
from ml_dataset import ProblemDataset   # 训练数据集

# ==================== 核心功能函数 ====================

def train_models(model_dir: str = "models", test_size: float = 0.1):
    """
    训练所有支持的机器学习模型
    
    这是脚本的核心功能，负责训练朴素贝叶斯、逻辑回归、随机森林三种模型，
    比较它们的性能，选择最佳模型并保存到指定目录。
    
    参数:
        model_dir (str, optional): 模型保存目录，默认"models"
        test_size (float, optional): 测试集比例，默认0.1（10%）
        
    返回:
        tuple: (model_path, results) 元组
            - model_path (str): 最佳模型的保存路径
            - results (dict): 所有模型的训练结果字典
            
    训练流程:
        1. 创建机器学习训练器实例
        2. 显示数据集统计信息（样本数、类别分布、标签分布）
        3. 训练所有支持的模型类型
        4. 比较模型性能并选择最佳模型
        5. 保存最佳模型到指定目录
        6. 显示最佳模型的性能指标
        
    性能指标:
        - 测试准确率：模型在测试集上的表现
        - 交叉验证：模型泛化能力的评估
        - 特征数：输入特征的维度
        - 标签数：分类类别的数量
        
    使用示例:
        # 使用默认参数训练
        model_path, results = train_models()
        
        # 自定义参数训练
        model_path, results = train_models("my_models", 0.2)
        
        # 查看训练结果
        for model_type, result in results.items():
            print(f"{model_type}: {result['test_accuracy']:.4f}")
    """
    print("开始训练机器学习模型...")
    print("=" * 60)
    
    # 创建机器学习训练器实例
    trainer = MLModelTrainer()
    
    # 显示数据集统计信息
    dataset = trainer.dataset
    print(f"数据集信息:")
    print(f"  总样本数: {len(dataset.training_data)}")
    print(f"  类别数: {len(dataset.get_categories())}")
    print(f"  标签数: {len(dataset.get_labels())}")
    
    # 显示类别分布统计
    print(f"\n类别分布:")
    categories = dataset.get_categories()
    for category in categories:
        count = sum(1 for item in dataset.training_data if item["category"] == category)
        print(f"  {category}: {count} 个样本")
    
    # 显示标签分布统计
    print(f"\n标签分布:")
    labels = dataset.get_labels()
    for label in labels:
        count = sum(1 for item in dataset.training_data if item["label"] == label)
        print(f"  {label}: {count} 个样本")
    
    print(f"\n{'='*60}")
    
    # 训练所有支持的模型类型
    # 包括朴素贝叶斯、逻辑回归、随机森林
    results = trainer.train_all_models(test_size)
    
    # 保存最佳模型到指定目录
    model_path = trainer.save_best_model(model_dir)
    
    print(f"\n{'='*60}")
    print("训练完成!")
    print(f"最佳模型已保存到: {model_path}")
    
    # 显示最佳模型的详细信息
    best_model_type = max(results.items(), key=lambda x: x[1]['test_accuracy'])[0]
    best_results = results[best_model_type]
    
    print(f"\n最佳模型性能:")
    print(f"  模型类型: {best_model_type}")
    print(f"  测试准确率: {best_results['test_accuracy']:.4f}")
    print(f"  交叉验证: {best_results['cv_mean']:.4f} ± {best_results['cv_std']:.4f}")
    print(f"  特征数: {best_results['feature_count']}")
    print(f"  标签数: {best_results['label_count']}")
    
    return model_path, results

def test_model(model_path: str):
    """
    测试训练好的机器学习模型
    
    使用预定义的测试用例验证模型的预测能力和准确性。
    测试用例涵盖了各种编程问题类型，用于评估模型的泛化能力。
    
    参数:
        model_path (str): 模型文件的路径
        
    功能:
        - 加载训练好的模型
        - 使用多种测试用例进行预测
        - 显示预测结果和置信度
        - 处理预测异常情况
        
    测试用例类型:
        - 数学运算：计算两个整数的和
        - 数组操作：找到数组中的最大值
        - 字符串处理：反转字符串
        - 数学问题：计算阶乘、判断质数
        - 特殊问题：计算煤球总数
        - 排序算法：快速排序
        - 图算法：深度优先搜索
        - 动态规划：背包问题
        - 递归算法：斐波那契数列
        
    输出格式:
        对于每个测试用例，显示：
        - 问题描述
        - 预测结果
        - 置信度分数
        
    异常处理:
        - 模型加载失败
        - 预测过程异常
        - 文件不存在
        
    使用示例:
        # 测试指定模型
        test_model("models/best_model_random_forest.joblib")
        
        # 测试自定义模型
        test_model("my_models/custom_model.joblib")
    """
    print(f"\n测试模型: {model_path}")
    print("=" * 60)
    
    # 创建训练器实例用于加载模型
    from ml_trainer import MLModelTrainer
    trainer = MLModelTrainer()
    
    # 加载训练好的最佳模型
    model = trainer.load_best_model(model_path)
    
    # 定义测试用例
    # 涵盖各种编程问题类型，用于验证模型的泛化能力
    test_cases = [
        "计算两个整数的和",                    # 数学运算
        "找到数组中的最大值",                  # 数组操作
        "反转字符串",                         # 字符串处理
        "计算阶乘",                           # 数学问题
        "判断一个数是否为质数",                # 数学判断
        "计算煤球的总数",                     # 特殊数学问题
        "使用快速排序对数组进行排序",          # 排序算法
        "实现深度优先搜索",                   # 图算法
        "解决背包问题",                       # 动态规划
        "计算斐波那契数列的第n项"             # 递归算法
    ]
    
    print("测试用例预测结果:")
    for i, test_case in enumerate(test_cases, 1):
        try:
            # 使用模型进行预测
            prediction, confidence = model.predict(test_case)
            print(f"{i:2d}. 问题: {test_case}")
            print(f"    预测: {prediction} (置信度: {confidence:.4f})")
        except Exception as e:
            # 处理预测异常
            print(f"{i:2d}. 问题: {test_case}")
            print(f"    错误: {e}")
        print()

# ==================== 命令行接口 ====================

def main():
    """
    主函数 - 命令行接口
    
    提供完整的命令行接口，支持模型训练和测试功能。
    使用argparse库解析命令行参数，提供灵活的操作选项。
    
    支持的命令:
        --train: 训练所有支持的机器学习模型
        --test: 测试指定的模型文件
        --model-dir: 指定模型保存目录
        --test-size: 指定测试集比例
        
    参数说明:
        --train: 布尔标志，触发模型训练流程
        --test: 字符串参数，指定要测试的模型文件路径
        --model-dir: 字符串参数，模型保存目录，默认"models"
        --test-size: 浮点数参数，测试集比例，默认0.1
        
    使用流程:
        1. 解析命令行参数
        2. 根据参数执行相应操作
        3. 提供操作结果和使用建议
        
    异常处理:
        - 模型文件不存在
        - 参数格式错误
        - 操作失败
        
    使用示例:
        # 训练模型
        python train_ml_models.py --train
        
        # 测试模型
        python train_ml_models.py --test models/best_model_random_forest.joblib
        
        # 自定义参数训练
        python train_ml_models.py --train --model-dir my_models --test-size 0.2
        
        # 查看帮助
        python train_ml_models.py --help
    """
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description="机器学习模型训练脚本")
    
    # 添加训练参数
    parser.add_argument("--train", action="store_true", help="训练模型")
    
    # 添加测试参数
    parser.add_argument("--test", type=str, help="测试指定模型")
    
    # 添加模型目录参数
    parser.add_argument("--model-dir", type=str, default="models", help="模型保存目录")
    
    # 添加测试集比例参数
    parser.add_argument("--test-size", type=float, default=0.1, help="测试集比例")
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 根据参数执行相应操作
    if args.train:
        # 执行模型训练
        model_path, results = train_models(args.model_dir, args.test_size)
        print(f"\n模型训练完成，可以运行以下命令测试模型:")
        print(f"python train_ml_models.py --test {model_path}")
    
    elif args.test:
        # 检查模型文件是否存在
        if not os.path.exists(args.test):
            print(f"错误: 模型文件不存在: {args.test}")
            return
        # 执行模型测试
        test_model(args.test)
    
    else:
        # 显示使用说明
        print("请指定操作:")
        print("  --train    训练模型")
        print("  --test     测试模型")
        print("\n示例:")
        print("  python train_ml_models.py --train")
        print("  python train_ml_models.py --test models/best_model_random_forest.joblib")

# ==================== 主程序入口 ====================

if __name__ == "__main__":
    """
    主程序入口
    
    当直接运行此脚本时，调用main函数启动命令行接口。
    这是脚本的入口点，负责处理命令行参数并执行相应操作。
    
    执行流程:
        1. 检查是否直接运行脚本（而非导入）
        2. 调用main函数处理命令行参数
        3. 根据参数执行训练或测试操作
        
    支持的运行方式:
        - 直接运行: python train_ml_models.py --train
        - 测试模型: python train_ml_models.py --test model.joblib
        - 查看帮助: python train_ml_models.py --help
        
    注意事项:
        - 确保src目录存在且包含必要的模块
        - 训练前确保有足够的数据集
        - 测试前确保模型文件存在且格式正确
    """
    main()
